from dataclasses import dataclass, field
from typing import (
    List,
    Optional,
    Set,
    Iterator,
)
from pathlib import Path
import ast

from .base import IRNode
from ..symbol_table import Symbol, SymbolTable, SymbolRegistry
from ..symbol_meta import SymbolKind, TypeInfo


@dataclass
class IRModule(IRNode):
    """IR 模块，一个 python 文件的 IR 集合"""

    node_type: str = "module"

    nodes: List[IRNode] = field(default_factory=list)

    # 保留的元数据
    includes: Set[str] = field(default_factory=set)
    module_full_name: Optional[str] = field(default=None)
    source_root: Optional[Path] = field(default=None)
    symbol_table: SymbolTable = field(default_factory=SymbolTable)
    module_path: Path = field(default=Path("."))

    def creates_scope(self) -> bool:
        """模块节点创建作用域"""
        return True

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        for node in self.nodes:
            yield node
        yield from super().get_child_nodes()

    def add_node(self, node: IRNode) -> None:
        """添加IR节点，保持源码顺序"""
        self.nodes.append(node)
        node.set_parent(self)

        # 如果添加的是 TypeVar 节点，清除缓存
        if hasattr(node, "node_type") and node.node_type == "typevar":
            self.set_extension_param("_typevars_cache", None)

    def insert_node_after(self, target_node: IRNode, new_node: IRNode) -> None:
        """在指定节点后插入新节点"""
        try:
            index = self.nodes.index(target_node)
            self.nodes.insert(index + 1, new_node)
        except ValueError:
            # 如果找不到目标节点，添加到末尾
            self.nodes.append(new_node)
        new_node.set_parent(self)

    def insert_node_before(self, target_node: IRNode, new_node: IRNode) -> None:
        """在指定节点前插入新节点"""
        try:
            index = self.nodes.index(target_node)
            self.nodes.insert(index, new_node)
        except ValueError:
            # 如果找不到目标节点，添加到开头
            self.nodes.insert(0, new_node)
        new_node.set_parent(self)

    def replace_node(self, target_node: IRNode, new_node: IRNode) -> None:
        """替换指定节点"""
        try:
            index = self.nodes.index(target_node)
            # 清除旧节点的父节点关系
            if target_node.parent == self:
                target_node.set_parent(None)
            self.nodes[index] = new_node
        except ValueError:
            # 如果找不到目标节点，添加到末尾
            self.nodes.append(new_node)
        new_node.set_parent(self)
        # 统一处理所有直接子节点的 parent 设置
        # 确保所有子节点的 parent 链都指向 new_node，这样 lookup_symbol 才能正确向上查找符号
        for child in new_node.get_child_nodes():
            child.set_parent(new_node)

    def get_nodes_by_type(self, node_type: str, node_class=None) -> list:
        """根据节点类型获取节点列表

        Args:
            node_type: 节点类型字符串
            node_class: 可选的节点类型，用于类型提示

        Returns:
            指定类型的节点列表
        """
        result = [node for node in self.nodes if node.node_type == node_type]
        return result

    def get_class_nodes(self) -> List[IRNode]:
        """获取所有类节点"""
        return [node for node in self.nodes if node.node_type == "class"]

    def get_function_nodes(self) -> List[IRNode]:
        """获取所有函数节点"""
        return [node for node in self.nodes if node.node_type == "function"]

    def get_enum_nodes(self) -> List[IRNode]:
        """获取所有枚举节点"""
        return [node for node in self.nodes if node.node_type == "enum"]

    def get_local_symbols(self) -> Iterator[Symbol]:
        """获取模块中定义的所有本地符号

        Returns:
            本地符号的迭代器
        """
        # 遍历所有节点，收集符号定义
        for node in self.nodes:
            # 只有可以成为符号的节点才调用 get_symbol
            if node.node_type in ["class", "function", "enum", "typevar"]:
                symbol = node.get_symbol()
                if symbol:
                    yield symbol
            # 对于赋值语句等特殊情况，仍然需要收集其本地符号
            elif node.node_type == "assign":
                from .statements.assignment import IRAssign

                if isinstance(node, IRAssign):
                    yield from node.get_local_symbols()
            elif node.node_type == "import":
                # 处理导入符号
                yield from node.get_local_symbols()

    def get_parent_scope(self) -> Optional[IRNode]:
        """获取父作用域节点

        Returns:
            模块没有父作用域，返回 None
        """
        return None

    def require_include(self, include_path: str) -> None:
        """要求包含指定头文件

        Args:
            include_path: 头文件路径，如 "<iostream>" 或 "\"local.h\""
        """
        if include_path:
            self.includes.add(include_path)

    def require_includes(self, include_paths: List[str]) -> None:
        """要求包含多个头文件

        Args:
            include_paths: 头文件路径列表
        """
        for include_path in include_paths:
            self.require_include(include_path)

    def get_includes(self) -> Set[str]:
        """获取所有需要的头文件

        Returns:
            头文件路径集合
        """
        return self.includes.copy()

    def clear_includes(self) -> None:
        """清除所有头文件依赖"""
        self.includes.clear()

    def get_module_typevars(self) -> set:
        """获取模块中定义的所有 TypeVar 名称（带缓存）

        根据 Python 规范，TypeVar 应该在模块级别定义，不应该在局部作用域定义。
        此方法只查找模块的直接子节点，不递归查找。

        Returns:
            set: TypeVar 名称集合
        """
        # 检查缓存
        cached = self.get_extension_param("_typevars_cache")
        if cached is not None:
            return cached

        # 构建 TypeVar 集合
        typevars = set()
        for node in self.get_child_nodes():
            # 只查找模块的直接子节点，不递归
            if hasattr(node, "node_type") and node.node_type == "typevar":
                typevars.add(node.name)

        # 缓存结果
        self.set_extension_param("_typevars_cache", typevars)
        return typevars

    def build_from_ast(self, node: ast.Module) -> None:
        """从 AST 模块节点构建 IR 模块内容

        Args:
            node: AST 模块节点
        """
        super().build_from_ast(node)

        # 处理模块体中的所有语句
        for stmt in node.body:
            ir_node = self._create_ir_node_from_statement(stmt)
            if ir_node:
                ir_node.build_from_ast(stmt)
                self.add_node(ir_node)

    def _create_ir_node_from_statement(self, stmt: ast.stmt) -> Optional[IRNode]:
        """根据 AST 语句创建对应的空 IR 节点

        Args:
            stmt: AST 语句节点

        Returns:
            创建的空 IR 节点，如果不支持则返回 None
        """
        from .declarations import IRClassNode, IRFunctionNode, IRImportNode
        from .statements import IRAssign, IRAugAssign, IRExprStmt

        if hasattr(stmt, "lineno"):
            source_line = stmt.lineno
        else:
            source_line = 0

        # 根据 AST 节点类型创建对应的空 IR 节点
        if hasattr(stmt, "__class__"):
            import ast

            if isinstance(stmt, ast.ClassDef):
                return IRClassNode(name=stmt.name, source_line=source_line, parent=self)
            elif isinstance(stmt, ast.FunctionDef) or isinstance(
                stmt, ast.AsyncFunctionDef
            ):
                return IRFunctionNode(
                    name=stmt.name, source_line=source_line, parent=self
                )
            elif isinstance(stmt, (ast.Import, ast.ImportFrom)):
                if self.is_skip_import(stmt):
                    return None
                # 对于导入语句，创建导入节点
                return IRImportNode(source_line=source_line, parent=self)
            elif isinstance(stmt, ast.Assign):
                # 检查是否是 TypeVar 赋值
                if self._is_typevar_assignment(stmt):
                    from .declarations.typevar import IRTypeVar

                    return IRTypeVar(source_line=source_line, parent=self)
                # 对于普通赋值语句，创建赋值节点
                return IRAssign(source_line=source_line, parent=self)
            elif isinstance(stmt, ast.AnnAssign):
                # 对于注解赋值语句，创建赋值节点
                return IRAssign(source_line=source_line, parent=self)
            elif isinstance(stmt, ast.AugAssign):
                # 对于增强赋值语句，创建增强赋值节点
                return IRAugAssign(source_line=source_line, parent=self)
            elif isinstance(stmt, ast.Expr):
                # 对于表达式语句，创建表达式语句节点
                return IRExprStmt(source_line=source_line, parent=self)
            elif isinstance(stmt, ast.Global):
                # 对于 global 语句，创建 global 节点
                from .statements.global_stmt import IRGlobal

                return IRGlobal(source_line=source_line, parent=self)
            elif isinstance(stmt, ast.Nonlocal):
                # 对于 nonlocal 语句，创建 nonlocal 节点
                from .statements.global_stmt import IRNonlocal

                return IRNonlocal(source_line=source_line, parent=self)
            # 其他语句类型暂时跳过

        return None

    def _is_typevar_assignment(self, stmt) -> bool:
        """检查是否是 TypeVar 赋值语句，如 T = TypeVar("T")"""
        import ast

        if not isinstance(stmt, ast.Assign) or len(stmt.targets) != 1:
            return False

        target = stmt.targets[0]
        if not isinstance(target, ast.Name):
            return False

        # 检查右侧是否是 TypeVar 调用
        if isinstance(stmt.value, ast.Call):
            if (
                isinstance(stmt.value.func, ast.Name)
                and stmt.value.func.id == "TypeVar"
            ):
                return True

        return False

    def is_skip_import(self, stmt: "ast.stmt") -> bool:
        """检查是否需要跳过导入语句"""
        if isinstance(stmt, ast.ImportFrom) and stmt.module == "typing":
            return True
        return False

    def lookup_symbol(
        self,
        name: str,
        owner_type: Optional[TypeInfo] = None,
        symbol_kinds: Optional[List[SymbolKind]] = None,
        *,
        template_args: Optional[List[TypeInfo]] = None,
        func_args: Optional[List[TypeInfo]] = None,
    ) -> Optional[Symbol]:
        """扩展的符号查找：本地符号 -> 导入符号 -> 内置符号"""

        # 1. 查找 IR 节点定义的符号
        symbol = super().lookup_symbol(
            name,
            owner_type=owner_type,
            symbol_kinds=symbol_kinds,
            func_args=func_args,
            template_args=template_args,
        )
        if symbol:
            return symbol

        # 2. 查找符号表中的符号
        symbol = self.symbol_table.lookup_symbol(
            name,
            template_args=template_args,
            func_args=func_args,
            symbol_kinds=symbol_kinds,
        )
        if symbol:
            return symbol

        # 3. 如果没有父模块，查找全局符号
        if self.parent is None:
            return self.symbol_table.lookup_global_symbol(
                name,
                template_args=template_args,
                func_args=func_args,
                symbol_kinds=symbol_kinds,
            )

        return None

    def get_registry(self) -> SymbolRegistry:
        """获取符号注册表"""
        return self.symbol_table.symbol_registry
